#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : SiamCNN_LSTM.py
# @Author: Richard Chiming Xu
# @Date  : 2022/1/17
# @Desc  :

'''
    1. 使用训练好的wv初始化embedding
    2. SiamCNN (Word2Vec句子编码 + 1D CNN +FC)
    3. SiamLSTM (Word2Vec句子编码 + LSTM + FC)
'''

import re
import jieba
import pandas as pd
import numpy as np
from gensim.models import Word2Vec
import pickle
from tqdm import tqdm


class Config():
    # 数据加载部分
    dataset = 'paws-x'
    build_vocab = False  # 是否重构词典
    build_wv = False  # 是否重构词向量
    load_with_words = True  # 加载带分词的数据
    vocab_path = 'vocab'  # 路径
    wv_model_path = 'wv_model'  # 词向量名
    seq_len = 27  # 句子长度
    # 模型部分
    model_type = 'RNN'
    embed_dim = 100  # 词向量或embedding维度
    vocab_size = 100  # 词典大小，读取词典后会更新
    update_embed = True  # embedding是否随网络训练更新
    load_model = False # 是否加载已有模型预测
    save_model = True # 是否保存训练好的模型
    # 训练部分
    device = 'cpu'
    learning_rate = 1e-3
    batch_size = 128  # batch大小
    epochs = 50  # 训练次数
    print_loss = 100  # 打印loss次数


# 读取数据
def get_stopwords():
    stop_words = []
    with open('dataprocess/baidu_stopwords.txt', 'r', encoding='utf-8') as f:
        for line in f.readlines():
            stop_words.append(line.replace('\n', ''))
    return stop_words


# jieba分词
def cut(content, stop_words, config):
    # 去除符号
    content = re.sub("[\s+\.\!\/_,$%^*(+\"\']+|[+——！，。？、~@#￥%……&*（）]", "", content)

    result = []

    seg_list = jieba.lcut(content, cut_all=True)
    for i in seg_list:
        if i not in stop_words:
            result.append(i)

    if len(result) < config.seq_len:  # 小于规定长度，填充
        new_result = ['PAD' for i in range(config.seq_len)]
        new_result[:len(result)] = result
        return new_result
    else:
        return result[:config.seq_len]


# 保存pickle
def dump_pickle(obj, file_path):
    with open(file_path, 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


# 加载pickle
def load_pickle(file_path):
    with open(file_path, 'rb') as f:
        return pickle.load(f)


# 分词并去除停用词
def data_anaysis(df, stop_words, config):
    # 分词
    df['words_a'] = df['text_a'].apply(lambda x: cut(x, stop_words, config))
    df['words_b'] = df['text_b'].apply(lambda x: cut(x, stop_words, config))
    return df


# 将中文词转换为词典数字
def word2num(content, vocab):
    result = []
    for word in content:
        result.append(vocab[word])
    return result


# 构建词典
def build_vocab(train, dev, test, config, file_name):
    word_index = 0
    vocab_dict = {}
    dataset = [train, test, dev]
    for data in tqdm(dataset, desc='构建词典'):
        for i in range(len(data)):
            row = data.iloc[i]
            words_a = row['words_a']
            words_b = row['words_b']
            for word in words_a:
                if word not in vocab_dict.keys():
                    vocab_dict[word] = word_index
                    word_index += 1
            for word in words_b:
                if word not in vocab_dict.keys():
                    vocab_dict[word] = word_index
                    word_index += 1
    # 保存词典
    dump_pickle(vocab_dict, config.vocab_path + '/vocab_' + file_name + '.pkl')
    return vocab_dict


def data_reader(file_name, config):
    print('加载数据...')
    if config.load_with_words is True:
        train = load_pickle('data/' + file_name + '_train.pkl')
        dev = load_pickle('data/' + file_name + '_dev.pkl')
        test = load_pickle('data/' + file_name + '_test.pkl')
    else:
        # 获取停用词
        stop_words = get_stopwords()
        train = pd.read_csv('data/' + file_name + '/train.tsv', sep='\t', names=['text_a', 'text_b', 'label'])
        dev = pd.read_csv('data/' + file_name + '/dev.tsv', sep='\t', names=['text_a', 'text_b', 'label'])
        test = pd.read_csv('data/' + file_name + '/test.tsv', sep='\t', names=['text_a', 'text_b'])

        if len(set(train['label'])) > 2:
            train = train[train['label'].isin(['0', '1'])]
            train['label'] = train['label'].astype('int')

        if len(set(train['label'])) > 2:
            dev = dev[dev['label'].isin(['0', '1'])]
            dev['label'] = dev['label'].astype('int')

        test['label'] = -1
        train = train.dropna()
        dev = dev.dropna()

        # 分词并去除停用词
        print('分词并去除停用词...')
        train = data_anaysis(train, stop_words, config)
        dev = data_anaysis(dev, stop_words, config)
        test = data_anaysis(test, stop_words, config)
        # 保存分词后的数据
        dump_pickle(train, 'data/' + file_name + '_train.pkl')
        dump_pickle(train, 'data/' + file_name + '_dev.pkl')
        dump_pickle(test, 'data/' + file_name + '_test.pkl')

    # 统计词典
    print('加载词典...')
    if config.build_vocab is True:
        vocab = build_vocab(train, dev, test, config, file_name)
    else:
        vocab = load_pickle(config.vocab_path + '/vocab_' + file_name + '.pkl')

    # 将每个词转为数字
    print('将每个词转为词典数字...')
    train['word2num_a'] = train['words_a'].apply(lambda x: word2num(x, vocab))
    train['word2num_b'] = train['words_b'].apply(lambda x: word2num(x, vocab))
    dev['word2num_a'] = dev['words_a'].apply(lambda x: word2num(x, vocab))
    dev['word2num_b'] = dev['words_b'].apply(lambda x: word2num(x, vocab))
    test['word2num_a'] = test['words_a'].apply(lambda x: word2num(x, vocab))
    test['word2num_b'] = test['words_b'].apply(lambda x: word2num(x, vocab))

    return train, dev, test, vocab


# 加载词向量
def load_wv(train, dev, test, config, file_name):
    if config.build_wv is True:
        print('构建词向量...')
        # 若不存在词向量，则重新训练训练词向量
        datalist = [train, dev, test]

        context = []
        for data in datalist:
            for i in range(len(data)):
                row = data.iloc[i]
                context.append(row['word2num_a'])
                context.append(row['word2num_b'])

        wv_model = Word2Vec(sentences=context, vector_size=config.embed_dim, window=5, min_count=1, workers=4)
        wv_model.train(context, total_examples=1, epochs=1)
        wv_model.save(config.wv_model_path + '/wv_' + file_name)
    else:
        print('加载词向量...')
        wv_model = Word2Vec.load(config.wv_model_path + '/wv_' + file_name)

    return wv_model


from torch.utils.data import DataLoader, Dataset


# 自定义数据集
class SimDataset(Dataset):
    def __init__(self, df):
        super(SimDataset, self).__init__()
        self.text_a = df['word2num_a']
        self.text_b = df['word2num_b']
        self.label = df['label']
        self.len = len(df)

    def __getitem__(self, idx):
        vector_a = np.array(self.text_a.iloc[idx], dtype='int64')
        vector_b = np.array(self.text_b.iloc[idx], dtype='int64')
        label = np.array(self.label.iloc[idx]).astype("int64")

        return {'vector_a': vector_a,
                'vector_b': vector_b,
                'label': label}

    def __len__(self):
        return self.len


from torch import nn
from torch import functional as F
import torch

class LinModel(nn.Module):

    def __init__(self, in_features, out_features):
        super(LinModel, self).__init__()


        self.fc_1 = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(),
            nn.Dropout(0.02)
        )
        self.fc_2 = nn.Sequential(
            nn.Linear(256, 32),
            nn.ReLU(),
            nn.Dropout(0.02)
        )
        self.fc_3 = nn.Sequential(
            nn.Linear(32, 4),
            nn.ReLU(),
            nn.Dropout(0.02)
        )
        self.fc_4 = nn.Sequential(
            nn.Linear(4, out_features),
        )
        self.softmax = nn.Softmax(1)

    def forward(self, X):
        X = self.fc_1(X)
        X = self.fc_2(X)
        X = self.fc_3(X)
        output = self.fc_4(X)

        return self.softmax(output)



class SiamCNN(nn.Module):
    def __init__(self, wv_mode, config):
        super(SiamCNN, self).__init__()
        self.device = config.device

        word_vectors = torch.randn([config.vocab_size, config.embed_dim])
        for i in range(0, config.vocab_size):
            word_vectors[i, :] = torch.from_numpy(wv_mode.wv[i])
        # 创建embedding层
        self.embedding = nn.Embedding.from_pretrained(word_vectors, freeze=config.update_embed)  # (32, 27, 100)
        if config.update_embed is False:
            self.embedding.weight.requires_grad = False

        self.conv_1 = nn.Sequential(
            nn.Conv1d(in_channels=config.seq_len, out_channels=16, kernel_size=2, stride=1),
            nn.ReLU(),
            nn.MaxPool1d(3))

        self.conv_2 = nn.Sequential(
            nn.Conv1d(in_channels=config.seq_len, out_channels=16, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.MaxPool1d(3))

        self.conv_3 = nn.Sequential(
            nn.Conv1d(in_channels=config.seq_len, out_channels=16, kernel_size=5, stride=1),
            nn.ReLU(),
            nn.MaxPool1d(3))

        self.flattern = nn.Flatten()

        # 定义池化层
        self.max_pool = nn.MaxPool1d(3)

        # 定义线性层
        self.lin_model = LinModel(1552, 2)


    # 计算两个向量的相似度
    def cos_sim(self, vector_a, vector_b):
        """
        计算两个向量之间的余弦相似度
        :param vector_a: 向量 a
        :param vector_b: 向量 b
        :return: sim
        """
        return torch.tensor([torch.cosine_similarity(vector_a, vector_b, 0, 1e-8)])

    def forward_one(self, text):
        # 计算句子A
        x = self.embedding(text)
        conv_1 = self.conv_1(x)
        conv_2 = self.conv_2(x)
        conv_3 = self.conv_3(x)
        # 合并各卷积结果取最大值
        x = torch.cat([conv_1, conv_2, conv_3], 2)
        x = x.view(x.size(0), -1)
        return self.lin_model(x)

    def forward(self, words_a, words_b):
        # words_a (batch_size, seq_len)(32, 27)
        # 计算句子A
        x_a = self.forward_one(words_a)
        # 计算句子B
        x_b = self.forward_one(words_b)
        return x_a, x_b


class SiamLSTM(nn.Module):
    def __init__(self, wv_mode, config):
        super(SiamLSTM, self).__init__()
        self.device = config.device

        word_vectors = torch.randn([config.vocab_size, config.embed_dim])
        for i in range(0, config.vocab_size):
            word_vectors[i, :] = torch.from_numpy(wv_mode.wv[i])
        # 创建embedding层
        self.embedding = nn.Embedding.from_pretrained(word_vectors, freeze=config.update_embed)  # (32, 27, 100)
        if config.update_embed is False:
            self.embedding.weight.requires_grad = False
        # 创建rnn
        self.rnn = nn.LSTM(input_size=config.embed_dim, hidden_size=10, num_layers=1)
        # 创建线性层
        self.lin_model = LinModel(270, 2)

    def forward_one(self, text):
        # 计算a
        x = self.embedding(text)  # embedding转换
        # rnn
        x = x.transpose(0, 1)  # 交换维度，因为RNN的输入是 （L, D, H）
        x, _ = self.rnn(x)
        x = x.transpose(0, 1)  # 还原维度，因为RNN的输出是 （L, D, H)
        x = x.contiguous().view(x.size(0), -1)
        return self.lin_model(x)

    def forward(self,words_a, words_b):
        # 计算a
        x_a = self.forward_one(words_a) # embedding转换
        # 计算b
        x_b = self.forward_one(words_b)
        return x_a, x_b


class ContrastiveLoss(torch.nn.Module):

    def __init__(self, margin=2.0):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin

    def forward(self, output1, output2, label):
        euclidean_distance = torch.pairwise_distance(output1, output2)
        loss_contrastive = torch.mean((1 - label) * torch.pow(euclidean_distance, 2) /2 +
        (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2)/2)

        y_pred = []
        for i in euclidean_distance:
            if i <0.5:
                y_pred.append(0)
            else:
                y_pred.append(1)

        return loss_contrastive, y_pred


from torch.optim import AdamW
from torch.optim.lr_scheduler import ExponentialLR
from sklearn.metrics import f1_score, accuracy_score

def predict(config, model, test_df):
    predict_labels = []
    # 创建dataloader
    test_dataset = SimDataset(test_df)
    test_dataloader = DataLoader(test_dataset, batch_size=config.batch_size)
    model.eval()
    with torch.no_grad():
        for mini_batch in test_dataloader:
            mini_batch = {item: value.to(config.device) for item, value in mini_batch.items()}
            # 获取数据
            text_a = mini_batch['vector_a']
            text_b = mini_batch['vector_b']
            x_a, x_b = model(text_a, text_b)
            euclidean_distance = torch.pairwise_distance(x_a, x_b)
            y_pred = []
            for i in euclidean_distance:
                if i < 0.5:
                    y_pred.append(0)
                else:
                    y_pred.append(1)
            indices = torch.tensor(y_pred, device=config.device)
            predict_labels += indices
    return torch.tensor(predict_labels).numpy()

# 校验
def evaluation(config, model, val_dataloader, loss_fn):
    model.eval()
    preds = []
    labels = []
    val_loss = 0.
    # val_iterator = tqdm(val_dataloader, desc='Evaluation', total=len(val_dataloader))
    with torch.no_grad():
        for mini_batch in val_dataloader:
            mini_batch = {item: value.to(config.device) for item, value in mini_batch.items()}
            # 获取数据
            text_a = mini_batch['vector_a']
            text_b = mini_batch['vector_b']
            label = mini_batch['label']
            labels += label

            x_a, x_b = model(text_a, text_b)

            loss, y_pred = loss_fn(x_a, x_b, label)
            y_pred = torch.tensor(y_pred, device=config.device)

            val_loss += loss.item()

            preds += y_pred

    avg_val_loss = val_loss / len(val_dataloader)
    labels = torch.tensor(labels).numpy()
    preds = torch.tensor(preds).numpy()
    f1 = f1_score(labels, preds, average='macro')
    # -----------new ----------------#
    acc = accuracy_score(labels, preds)
    # -----------new ----------------#
    return avg_val_loss, f1, acc


#  训练数据
def train(train_df, dev_df, wv_model, config):
    # 创建数据集
    train_dataset = SimDataset(train_df)
    dev_dataset = SimDataset(dev_df)
    train_dataloader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
    dev_dataloader = DataLoader(dev_dataset, batch_size=config.batch_size)

    # 创建模型
    if config.model_type == 'CNN':
        model = SiamCNN(wv_model, config)
    elif config.model_type == 'RNN':
        model = SiamLSTM(wv_model, config)
    else:
        raise Exception('错误模型类型！')

    # model = SiamCNN(wv_model, config, each_filter_num=128, filter_heights=[2, 3, 5])
    model.to(config.device)
    # 定义优化器
    opt = AdamW(lr=config.learning_rate, params=model.parameters())
    # scheduler = ExponentialLR(opt, gamma=0.9)
    # 定义损失函数
    loss_fn = ContrastiveLoss(1)

    # 遍历epoch，开始训练
    for epoch in range(config.epochs):
        model.train()
        for iter_id, mini_batch in enumerate(train_dataloader):
            mini_batch = {item: value.to(config.device) for item, value in mini_batch.items()}
            # 获取数据
            text_a = mini_batch['vector_a']
            text_b = mini_batch['vector_b']
            label = mini_batch['label']
            # 训练
            x_a, x_b = model(text_a, text_b)
            # 计算损失
            loss, y_pred = loss_fn(x_a, x_b, label)

            # 打印loss
            indices = torch.tensor(y_pred, device=config.device)
            correct = torch.sum(indices == label)

            avg_loss = torch.mean(loss)
            if iter_id % config.print_loss == 0:
                print('epoch: {}, iter: {}, loss is: {}, acc is: {}'.format(epoch, iter_id, avg_loss,
                                                                            correct.item() * 1.0 / len(text_a)))
            # 反向传播
            avg_loss.backward()
            # 最小化loss,更新参数
            opt.step()
            # scheduler.step()
            # 清除梯度
            model.zero_grad()

        # 完成1个epoch，验证
        avg_val_loss, f1, acc = evaluation(config, model, dev_dataloader, loss_fn)
        print('-'*50)
        print('epoch: {}, val_loss: {}, val_f1: {}, val_acc: {}'.format(epoch, avg_val_loss, f1,acc))
        print('-' * 50)

    return model

data_list = ['bq_corpus','lcqmc','paws-x']

for data in data_list:
    # 配置文件
    conf = Config()
    conf.device = 'cuda' if torch.cuda.is_available() else 'cpu'
    conf.dataset = data
    # 读取数据
    train_df, dev_df, test_df, vocab = data_reader(conf.dataset, conf)


    if conf.load_model is False:
        # 加载测向量
        wv_model = load_wv(train_df, dev_df, test_df, conf, conf.dataset)
        # 更新词典大小
        conf.vocab_size = len(vocab)
        # 训练模型
        model = train(train_df, dev_df, wv_model, conf)
        if conf.save_model is True:
            torch.save(model, 'model/'+conf.dataset+'.pkl')
    else:
        model = torch.load('model/'+conf.dataset+'.pkl')


    # 推理模型
    predict_labels = predict(conf, model, test_df)
    # 保存结果
    test_df['index'] = test_df.index
    test_df['prediction'] = predict_labels
    test_df.to_csv('result/'+conf.dataset+'.tsv', index=False, columns=['index','prediction'], sep='\t')



